package com.ctvit.breadking.news.core.utils;

import org.slf4j.MDC;

import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.function.Supplier;

import static com.ctvit.breadking.news.core.constant.CommonConstant.TRACE_ID;

/**
 * 线程MDC工具类
 *
 * @author heyingcheng
 * @date 2024/03/27
 */
public class ThreadMdcUtils {

    private ThreadMdcUtils() {
        throw new IllegalStateException("Utility class");
    }

    public static void setTraceIdIfAbsent() {
        String traceId = MDC.get(TRACE_ID);
        if (Objects.isNull(traceId)) {
            MDC.put(TRACE_ID, UUID.randomUUID().toString());
        } else {
            MDC.put(TRACE_ID, traceId + "-" + UUID.randomUUID().toString());
        }
    }

    /**
     * 封装给定的Callable对象，将其与MDC（Mapped Diagnostic Contexts）的当前上下文映射关联起来。
     *
     * @param callable 要被封装的Callable对象，它定义了一个带返回值的任务。
     * @param <T>      Callable返回的类型参数。
     * @return 包装后的Callable对象，当执行任务时，会将MDC的当前上下文映射复制到执行上下文中。
     * @see MDC#getCopyOfContextMap() 用于获取MDC的当前上下文映射副本的方法。
     */
    public static <T> Callable<T> wrapCallable(final Callable<T> callable) {
        return wrap(callable, MDC.getCopyOfContextMap());
    }

    private static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return () -> {
            if (Objects.isNull(context)) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                return callable.call();
            } finally {
                MDC.clear();
            }
        };
    }

    /**
     * 封装一个Runnable对象，将其与MDC（Mapped Diagnostic Context）的上下文映射关联起来。
     *
     * @param runnable 要被封装的Runnable对象。
     * @return 封装后的Runnable对象，当执行时，会带有MDC的上下文信息。
     * @see MDC#getCopyOfContextMap() 获取MDC当前的上下文映射副本
     */
    public static Runnable wrap(final Runnable runnable) {
        return wrap(runnable, MDC.getCopyOfContextMap());
    }

    private static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            if (Objects.isNull(context)) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }

    /**
     * 封装给定的Supplier对象，将其与MDC（Mapped Diagnostic Contexts）的当前上下文映射关联起来。
     *
     * @param supplier 要被封装的Callable对象，它定义了一个带返回值的任务。
     * @param <T>      Callable返回的类型参数。
     * @return 包装后的Callable对象，当执行任务时，会将MDC的当前上下文映射复制到执行上下文中。
     * @see MDC#getCopyOfContextMap() 用于获取MDC的当前上下文映射副本的方法。
     */
    public static <T> Supplier<T> wrapSupplier(final Supplier<T> supplier) {
        return wrap(supplier, MDC.getCopyOfContextMap());
    }

    private static <T> Supplier<T> wrap(final Supplier<T> supplier, final Map<String, String> context) {
        return () -> {
            if (Objects.isNull(context)) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                return supplier.get();
            } finally {
                MDC.clear();
            }
        };
    }

    /**
     * 执行给定的Runnable对象。
     * 该方法在执行Runnable之前和之后会清理MDC（Mapped Diagnostic Context）日志上下文，
     * 并在执行前设置跟踪ID（如果尚未设置），以确保日志的一致性和准确性。
     *
     * @param runnable 要执行的Runnable对象，不能为空。
     */
    public static void run(final Runnable runnable) {
        MDC.clear();
        setTraceIdIfAbsent();
        try {
            runnable.run();
        } finally {
            MDC.clear();
        }
    }

    /**
     * 通过给定的Supplier获取一个值。
     * 在获取值的过程中，会进行日志上下文的初始化和清除工作，确保日志追踪的准确性。
     *
     * @param supplier 一个提供所需值的Supplier实例，该Supplier的具体类型参数化为<T>。
     * @param <T>      返回值的类型参数。
     * @return 返回Supplier提供的值，其类型为T。
     */
    public static <T> T get(final Supplier<T> supplier) {
        MDC.clear();
        setTraceIdIfAbsent();
        try {
            return supplier.get();
        } finally {
            MDC.clear();
        }
    }

}
